package com.datastructures2.字符串.单词查找树;

import edu.princeton.cs.algs4.Queue;

/**
 * 基于三向单词查找树的符号表
 *
 * @author MaoLin Wang
 * @date 2020/2/2615:39
 */
public class TST<Value> {
    private Node root;

    private class Node<Value> {
        private char c;
        private Node left, mid, right;
        private Value value;
    }

    public Value get(String key) {
        Node x = get(root, key, 0);
        if (x == null) {
            return null;
        }
        return (Value) x.value;
    }

    private Node get(Node x, String key, int d) {
        if (x == null) {
            return null;
        }
        char c = key.charAt(d);
        if (c < x.c) {
            return get(x.left, key, d);
        } else if (c > x.c) {
            return get(x.right, key, d);
        } else if (d < key.length() - 1) {
            return get(x.mid, key, d + 1);
        } else {
            return x;
        }
    }

    public void put(String key, Value value) {
        root = put(root, key, value, 0);
    }

    private Node put(Node root, String key, Value value, int d) {
        char c = key.charAt(d);
        if (root == null) {
            root = new Node();
            root.c = c;
        }
        if (c < root.c) {
            root.left = put(root.left, key, value, d);
        } else if (c > root.c) {
            root.right = put(root.right, key, value, d);
        } else if (d < key.length() - 1) {
            root.mid = put(root.mid, key, value, d + 1);
        } else {
            root.value = value;
        }
        return root;

    }

    public String longestPrefixOf(String query) {
        if (query== null){
            throw new IllegalArgumentException("参数错误");
        }
        if (query.length()==0){
            return null;
        }
        int length=0;
        Node<Value>node=root;
        int d=0;
        while (node!=null && d< query.length()){
            char c=query.charAt(d);
            if (c<node.c){
                node=node.left;
            }else if (c>node.c){
                node=node.right;
            }else {
                d++;
                if (node.value!=null){
                    length=d;
                }
                node=node.mid;
            }
        }
        return query.substring(0,length);
    }
    public Iterable<String> keysWithPrefix(String prefix) {
        if (prefix == null) {
            throw new IllegalArgumentException("calls keysWithPrefix() with null argument");
        }
        Queue<String> queue = new Queue<String>();
        Node<Value> x = get(root, prefix, 0);
        if (x == null) return queue;
        if (x.value != null) queue.enqueue(prefix);
        collect(x.mid, new StringBuilder(prefix), queue);
        return queue;
    }

    // all keys in subtrie rooted at x with given prefix
    private void collect(Node<Value> x, StringBuilder prefix, Queue<String> queue) {
        if (x == null) return;
        collect(x.left,  prefix, queue);
        if (x.value != null) queue.enqueue(prefix.toString() + x.c);
        collect(x.mid,   prefix.append(x.c), queue);
        prefix.deleteCharAt(prefix.length() - 1);
        collect(x.right, prefix, queue);
    }
    public static void main(String[] args) {
        TrieST<Integer> trieST=new TrieST<>();
        trieST.put("shells",3);
        trieST.put("she",0);
        trieST.put("sea",2);
        trieST.put("sells",1);
      String shella = trieST.longstPrefixOf("shella");
        System.out.println(shella);


    }
}
